from fastapi import FastAPI, WebSocket, Request, WebSocketDisconnect
# FastAPI - 主框架类
# WebSocket - WebSocket处理类
# Request - HTTP请求处理类
# WebSocketDisconnect - WebSocket断开连接异常

from fastapi.responses import HTMLResponse
# HTMLResponse - 返回HTML响应的类

from fastapi.staticfiles import StaticFiles
# StaticFiles - 静态文件服务类

from fastapi.templating import Jinja2Templates
# Jinja2Templates - 模板渲染类

import uvicorn
# uvicorn - ASGI服务器实现

import json
# json - JSON处理模块

import asyncio
# asyncio - 异步IO库

import os
# os - 操作系统接口模块

from datetime import datetime
# datetime - 日期时间处理类

app = FastAPI()
# 创建FastAPI应用实例

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
# 获取当前文件的绝对路径的目录部分

app.mount("/static", StaticFiles(directory=os.path.join(BASE_DIR, "static")), name="static")
# 挂载静态文件服务
# 参数:
#   "/static" - URL路径前缀
#   directory - 静态文件目录路径
#   name="static" - 路由名称

templates = Jinja2Templates(directory=os.path.join(BASE_DIR, "templates"))
# 初始化模板引擎
# 参数:
#   directory - 模板文件目录路径

# 存储连接和用户
active_connections = {}
# 存储活跃的WebSocket连接，格式: {user_id: websocket}

online_users = {}
# 存储在线用户信息，格式: {user_id: user_info}
@app.get("/")
async def get(request: Request):
    # 处理根路径GET请求
    # 参数:
    #   request: Request - 包含HTTP请求信息的对象
    return templates.TemplateResponse("index.html", {"request": request})
    # 返回渲染后的模板
    # 参数:
    #   "index.html" - 模板文件名
    #   {"request": request} - 传递给模板的上下文

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # WebSocket端点处理
    # 参数:
    #   websocket: WebSocket - WebSocket连接对象
    
    await websocket.accept()
    # 接受WebSocket连接
    
    user_id = None
    # 初始化用户ID
    
    try:
        # 1. 等待注册消息
        data = await websocket.receive_text()
        # 接收文本消息
        print("data(注册时的消息):",data)
        message = json.loads(data)
        print("message(注册时的消息):",message)
        # 解析JSON消息
        print("message.get('type')(注册时的消息):",message.get("type"))
        if message.get("type") != "register":
            # 检查消息类型是否为注册
            await websocket.close(code=1003)
            # 关闭连接
            # 参数:
            #   code=1003 - WebSocket关闭代码(表示无效数据)
            return
            
        # 2. 用户注册
        user_id = message["user_id"]
        print("user_id:",user_id)
        # 获取用户ID
        
        if user_id in active_connections:
            # 检查用户是否已连接
            await websocket.close(code=1008)
            # 关闭连接
            # 参数:
            #   code=1008 - WebSocket关闭代码(表示冲突)
            return
            
        active_connections[user_id] = websocket
        print("active_connections[user_id]和websocket:",active_connections[user_id])
        # 存储WebSocket连接
        
        online_users[user_id] = {
            "id": user_id,
            "name": message.get("name", user_id),
            "avatar": message.get("avatar", f"https://api.dicebear.com/7.x/miniavs/svg?seed={user_id}"),
            "status": "online"
        }
        print("online_users[user_id]:",online_users[user_id])
        # 存储用户信息
        
        # 3. 广播更新后的用户列表
        await broadcast_users()
        # 调用广播用户列表函数
        
        # 4. 发送欢迎消息
        await websocket.send_text(json.dumps({
            "type": "system",
            "content": "Welcome to the chat!"
        }))
        # 发送欢迎消息
        
        # 5. 主消息循环
        while True:
            data = await websocket.receive_text()
            # 持续接收消息
            print("data（聊天时的消息）:",data)
            message = json.loads(data)
            # 解析消息
            print("message（聊天时的消息）:",message)
            print("message.get('type')（聊天时的消息）:",message.get("type"))
            if message["type"] == "message":
                # 处理聊天消息
                print("message[\'to\']（聊天时的消息）:",message["to"])
                target_id = message["to"]
                # 获取目标用户ID
                print("active_connections:",active_connections)
                if target_id in active_connections:
                    # 检查目标用户是否在线
                    await active_connections[target_id].send_text(json.dumps({
                        "type": "message",
                        "from": user_id,
                        "to": target_id,
                        "content": message["content"],
                        "timestamp": message.get("timestamp", datetime.now().isoformat())
                    }))
                    # 转发消息给目标用户
                    
    except WebSocketDisconnect:
        # 处理WebSocket断开连接异常
        print(f"User {user_id} disconnected")
        await handle_disconnect(user_id)
        # 调用断开连接处理函数
        
    except Exception as e:
        # 处理其他异常
        print(f"Error: {e}")
        await handle_disconnect(user_id)

async def handle_disconnect(user_id: str):
    # 处理用户断开连接
    # 参数:
    #   user_id: str - 用户ID
    
    if user_id in active_connections:
        del active_connections[user_id]
        # 从活跃连接中移除
        
    if user_id in online_users:
        del online_users[user_id]
        # 从在线用户中移除
        
    await broadcast_users()
    # 广播更新后的用户列表

async def broadcast_users():
    # 广播用户列表给所有连接的用户
    
    users = list(online_users.values())
    # 获取所有在线用户信息
    print("users:",users)
    message = json.dumps({
        "type": "users",
        "users": users
    })
    # 构建用户列表消息
    
    for user_id, ws in active_connections.items():
        # 遍历所有活跃连接
        try:
            await ws.send_text(message)
            # 发送用户列表
        except:
            continue

if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
    # 启动Uvicorn服务器
    # 参数:
    #   "main:app" - 要运行的应用(模块名:应用变量名)
    #   host="0.0.0.0" - 监听所有网络接口
    #   port=8000 - 监听端口
    #   reload=True - 启用代码热重载